#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
平安保险文档 Elasticsearch 搜索系统
功能：
1. 连接到Elasticsearch
2. 创建保险文档索引
3. 索引docs文件夹中的所有文档
4. 执行智能搜索
5. 显示搜索结果
"""

import os
import json
import logging
from datetime import datetime
from typing import List, Dict, Any
from pathlib import Path

# 导入颜色高亮工具
try:
    from color_highlight_utils import process_highlight_fragment, print_colored, highlight_text
except ImportError:
    # 如果导入失败，提供简单的替代函数
    def process_highlight_fragment(fragment, style='red_bold'):
        return fragment.replace('<HIGHLIGHT>', '[').replace('</HIGHLIGHT>', ']')
    
    def print_colored(text, color='white', style=None):
        print(text)
    
    def highlight_text(text, style='red_bold'):
        return f"[{text}]"
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module='pypdf')#忽略pypdf模块的警告信息


try:
    from elasticsearch import Elasticsearch
    from elasticsearch.helpers import bulk
except ImportError:
    print("请安装 elasticsearch 库: pip install elasticsearch")
    exit(1)

try:
    import PyPDF2
except ImportError:
    print("请安装 PyPDF2 库: pip install PyPDF2")
    exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class InsuranceDocumentSearcher:
    """平安保险文档搜索系统"""
    
    def __init__(self, es_host="localhost", es_port=9200, index_name="insurance_docs"):
        """
        初始化搜索系统
        
        Args:
            es_host: Elasticsearch主机地址
            es_port: Elasticsearch端口
            index_name: 索引名称
        """
        self.es_host = es_host
        self.es_port = es_port
        self.index_name = index_name
        self.es_client = None
        self.docs_folder = "docs"
        
    def connect_elasticsearch(self) -> bool:
        """
        连接到Elasticsearch
        
        Returns:
            bool: 连接是否成功
        """
        try:
            # 尝试使用配置文件
            try:
                from elasticsearch_config import get_elasticsearch_client
                self.es_client = get_elasticsearch_client()
                logger.info("使用elasticsearch_config.py中的配置")
            except ImportError:
                # 回退到基础配置 - 修复连接方式
                self.es_client = Elasticsearch(
                    "https://localhost:9200",
                    request_timeout=30,
                    # 如果需要认证，取消下面的注释并填入正确的用户名密码
                    basic_auth=("elastic", "ncUg27UP5aMSzSElVXAX"),
                    verify_certs=False,
                )
                logger.info("使用默认配置连接")
            
            # 测试连接
            if self.es_client.ping():
                logger.info(f"✅ 成功连接到Elasticsearch: {self.es_host}:{self.es_port}")
                
                # 显示集群信息
                cluster_info = self.es_client.info()
                logger.info(f"集群名称: {cluster_info['cluster_name']}")
                logger.info(f"ES版本: {cluster_info['version']['number']}")
                return True
            else:
                logger.error("❌ 无法连接到Elasticsearch")
                return False
                
        except Exception as e:
            logger.error(f"❌ 连接Elasticsearch失败: {str(e)}")
            logger.error("请确保Elasticsearch服务正在运行")
            return False
    
    def create_index(self) -> bool:
        """
        创建保险文档索引
        
        Returns:
            bool: 创建是否成功
        """
        try:
            # 定义索引映射 - 简化版本，不使用IK分词器
            index_mapping = {
                "mappings": {
                    "properties": {
                        "title": {
                            "type": "text"
                        },
                        "content": {
                            "type": "text"
                        },
                        "file_name": {
                            "type": "keyword"
                        },
                        "file_type": {
                            "type": "keyword"
                        },
                        "file_size": {
                            "type": "long"
                        },
                        "created_time": {
                            "type": "date"
                        },
                        "insurance_type": {
                            "type": "keyword"
                        },
                        "keywords": {
                            "type": "keyword"
                        }
                    }
                },
                "settings": {
                    "number_of_shards": 1,
                    "number_of_replicas": 0
                }
            }
            
            # 检查索引是否已存在
            if self.es_client.indices.exists(index=self.index_name):
                logger.info(f"📋 索引 '{self.index_name}' 已存在，删除旧索引...")
                self.es_client.indices.delete(index=self.index_name)
            
            # 创建新索引
            self.es_client.indices.create(
                index=self.index_name,
                body=index_mapping
            )
            
            logger.info(f"✅ 成功创建索引: {self.index_name}")
            return True
            
        except Exception as e:
            logger.error(f"❌ 创建索引失败: {str(e)}")
            return False
    
    def extract_text_from_pdf(self, pdf_path: str) -> str:
        """
        从PDF文件提取文本
        
        Args:
            pdf_path: PDF文件路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            with open(pdf_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                text = ""
                
                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    text += page.extract_text() + "\n"
                
                return text.strip()
                
        except Exception as e:
            logger.error(f"❌ 提取PDF文本失败 {pdf_path}: {str(e)}")
            return ""
    
    def read_text_file(self, file_path: str) -> str:
        """
        读取文本文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            str: 文件内容
        """
        try:
            encodings = ['utf-8', 'gbk', 'gb2312', 'utf-16']
            
            for encoding in encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as file:
                        return file.read()
                except UnicodeDecodeError:
                    continue
            
            logger.error(f"❌ 无法读取文件 {file_path}: 编码不支持")
            return ""
            
        except Exception as e:
            logger.error(f"❌ 读取文件失败 {file_path}: {str(e)}")
            return ""
    
    def extract_insurance_type(self, file_name: str, content: str) -> str:
        """
        从文件名和内容中提取保险类型
        
        Args:
            file_name: 文件名
            content: 文件内容
            
        Returns:
            str: 保险类型
        """
        insurance_types = {
            "雇主": "雇主责任险",
            "责任险": "责任保险",
            "意外险": "意外伤害保险",
            "财产": "财产保险",
            "装修": "装修保险",
            "施工": "施工保险",
            "交通": "交通意外险",
            "医疗": "医疗保险"
        }
        
        file_name_lower = file_name.lower()
        content_sample = content[:500].lower()
        
        for keyword, insurance_type in insurance_types.items():
            if keyword in file_name_lower or keyword in content_sample:
                return insurance_type
        
        return "综合保险"
    
    def extract_keywords(self, content: str) -> List[str]:
        """
        从内容中提取关键词
        
        Args:
            content: 文档内容
            
        Returns:
            List[str]: 关键词列表
        """
        keywords = []
        
        # 保险相关关键词
        insurance_keywords = [
            "保险", "保障", "理赔", "赔偿", "责任", "意外", "工伤", 
            "雇主", "员工", "职业", "疾病", "伤残", "死亡", "医疗",
            "财产", "损失", "风险", "保费", "保额", "条款", "承保"
        ]
        
        content_lower = content.lower()
        for keyword in insurance_keywords:
            if keyword in content_lower:
                keywords.append(keyword)
        
        return keywords[:10]  # 限制关键词数量
    
    def prepare_documents(self) -> List[Dict[str, Any]]:
        """
        准备要索引的文档
        
        Returns:
            List[Dict]: 文档列表
        """
        documents = []
        docs_path = Path(self.docs_folder)
        
        if not docs_path.exists():
            logger.error(f"❌ 文档文件夹不存在: {self.docs_folder}")
            return documents
        
        logger.info(f"📂 开始处理文档文件夹: {self.docs_folder}")
        
        for file_path in docs_path.iterdir():
            if file_path.is_file():
                try:
                    file_name = file_path.name
                    file_size = file_path.stat().st_size
                    file_type = file_path.suffix.lower()
                    
                    logger.info(f"📄 处理文件: {file_name}")
                    
                    # 提取文本内容
                    if file_type == '.pdf':
                        content = self.extract_text_from_pdf(str(file_path))
                    elif file_type in ['.txt', '.md']:
                        content = self.read_text_file(str(file_path))
                    else:
                        logger.warning(f"⚠️ 跳过不支持的文件类型: {file_name}")
                        continue
                    
                    if not content.strip():
                        logger.warning(f"⚠️ 文件内容为空: {file_name}")
                        continue
                    
                    # 提取文档信息
                    insurance_type = self.extract_insurance_type(file_name, content)
                    keywords = self.extract_keywords(content)
                    
                    # 创建文档
                    doc = {
                        "_index": self.index_name,
                        "_source": {
                            "title": file_name.replace(file_type, ""),
                            "content": content,
                            "file_name": file_name,
                            "file_type": file_type,
                            "file_size": file_size,
                            "created_time": datetime.now().isoformat(),
                            "insurance_type": insurance_type,
                            "keywords": keywords
                        }
                    }
                    
                    documents.append(doc)
                    logger.info(f"✅ 文档准备完成: {file_name} ({len(content)} 字符)")
                    
                except Exception as e:
                    logger.error(f"❌ 处理文件失败 {file_path}: {str(e)}")
                    continue
        
        logger.info(f"📊 总共准备了 {len(documents)} 个文档")
        return documents
    
    def index_documents(self, documents: List[Dict[str, Any]]) -> bool:
        """
        批量索引文档
        
        Args:
            documents: 文档列表
            
        Returns:
            bool: 索引是否成功
        """
        try:
            if not documents:
                logger.warning("⚠️ 没有文档需要索引")
                return False
            
            logger.info(f"📤 开始批量索引 {len(documents)} 个文档...")
            
            # 执行批量索引
            success_count, failed_items = bulk(
                self.es_client,
                documents,
                index=self.index_name,
                chunk_size=100,
                request_timeout=60
            )
            
            logger.info(f"✅ 成功索引 {success_count} 个文档")
            
            if failed_items:
                logger.warning(f"⚠️ {len(failed_items)} 个文档索引失败")
                for item in failed_items:
                    logger.error(f"失败项: {item}")
            
            # 刷新索引
            self.es_client.indices.refresh(index=self.index_name)
            logger.info("🔄 索引刷新完成")
            
            return True
            
        except Exception as e:
            logger.error(f"❌ 批量索引失败: {str(e)}")
            return False
    
    def search_documents(self, query: str, size: int = 10) -> Dict[str, Any]:
        """
        搜索文档
        
        Args:
            query: 搜索查询
            size: 返回结果数量
            
        Returns:
            Dict: 搜索结果
        """
        try:
            # 构建搜索查询
            search_body = {
                "query": {
                    "bool": {
                        "should": [
                            {
                                "multi_match": {
                                    "query": query,
                                    "fields": ["title^3", "content^2", "keywords^2"],
                                    "type": "best_fields",
                                    "fuzziness": "AUTO"
                                }
                            },
                            {
                                "match_phrase": {
                                    "content": {
                                        "query": query,
                                        "boost": 2
                                    }
                                }
                            }
                        ],
                        "minimum_should_match": 1
                    }
                },
                "highlight": {
                    "fields": {
                        "content": {
                            "fragment_size": 150,
                            "number_of_fragments": 3,
                            "pre_tags": ["<HIGHLIGHT>"],
                            "post_tags": ["</HIGHLIGHT>"]
                        }
                    }
                },
                "size": size,
                "_source": ["title", "file_name", "insurance_type", "keywords"]
            }
            
            logger.info(f"🔍 执行搜索查询: {query}")
            
            # 执行搜索
            response = self.es_client.search(
                index=self.index_name,
                body=search_body
            )
            
            return response
            
        except Exception as e:
            logger.error(f"❌ 搜索失败: {str(e)}")
            return {}
    
    def display_search_results(self, response: Dict[str, Any], query: str):
        """
        显示搜索结果
        
        Args:
            response: 搜索响应
            query: 搜索查询
        """
        if not response or 'hits' not in response:
            print("❌ 搜索结果为空")
            return
        
        hits = response['hits']
        total_hits = hits['total']['value']
        max_score = hits['max_score']
        
        print("\n" + "="*80)
        print(f"🔍 搜索查询: {query}")
        print(f"📊 找到 {total_hits} 个相关文档 (最高分: {max_score:.4f})")
        print("="*80)
        
        for i, hit in enumerate(hits['hits'], 1):
            source = hit['_source']
            score = hit['_score']
            
            print(f"\n📄 结果 {i} (相关度: {score:.4f})")
            print(f"标题: {source['title']}")
            print(f"文件: {source['file_name']}")
            print(f"保险类型: {source['insurance_type']}")
            print(f"关键词: {', '.join(source['keywords'])}")
            
            # 显示高亮片段
            if 'highlight' in hit and 'content' in hit['highlight']:
                print("📝 相关内容:")
                for fragment in hit['highlight']['content']:
                    # 使用颜色高亮工具处理片段
                    highlighted_fragment = process_highlight_fragment(fragment, 'red_bold')
                    print(f"   {highlighted_fragment}")
            
            print("-" * 60)
        
        print(f"\n✅ 搜索完成，共显示 {len(hits['hits'])} 个结果")


def main():
    """主函数"""
    print("🚀 启动平安保险文档搜索系统")
    print("="*60)
    
    # 创建搜索系统实例
    searcher = InsuranceDocumentSearcher(
        es_host="localhost",
        es_port=9200,
        index_name="insurance_docs"
    )
    
    # 1. 连接到Elasticsearch
    print("\n📡 步骤1: 连接到Elasticsearch")
    if not searcher.connect_elasticsearch():
        print("❌ 连接失败，程序退出")
        return
    
    # 2. 创建索引
    print("\n📋 步骤2: 创建保险文档索引")
    if not searcher.create_index():
        print("❌ 创建索引失败，程序退出")
        return
    
    # 3. 准备并索引文档
    print("\n📚 步骤3: 索引文档")
    documents = searcher.prepare_documents()
    if not searcher.index_documents(documents):
        print("❌ 索引文档失败，程序退出")
        return
    
    # 4. 执行搜索
    print("\n🔍 步骤4: 执行搜索")
    search_query = "工伤保险和雇主险有什么区别？"
    
    response = searcher.search_documents(search_query, size=5)
    
    # 5. 显示搜索结果
    print("\n📊 步骤5: 显示搜索结果")
    searcher.display_search_results(response, search_query)
    
    # 交互式搜索
    print("\n" + "="*60)
    print("💡 进入交互式搜索模式 (输入 'quit' 退出)")
    print("="*60)
    
    while True:
        try:
            user_query = input("\n🔍 请输入搜索查询: ").strip()
            
            if user_query.lower() in ['quit', 'exit', '退出']:
                print("👋 再见！")
                break
            
            if not user_query:
                print("⚠️ 请输入有效的搜索查询")
                continue
            
            response = searcher.search_documents(user_query, size=3)
            searcher.display_search_results(response, user_query)
            
        except KeyboardInterrupt:
            print("\n👋 程序被用户中断，再见！")
            break
        except Exception as e:
            print(f"❌ 搜索过程中出错: {str(e)}")


if __name__ == "__main__":
    main()